package cn.oldsix.spark.utils

import kafka.common.TopicAndPartition
import kafka.utils.ZkUtils
import org.I0Itec.zkclient.ZkClient
import org.apache.log4j.LogManager
import org.apache.spark.rdd.RDD
import org.apache.spark.streaming.kafka.{HasOffsetRanges, OffsetRange}

/*
* @Author : Wu.D.J
* @Create : 2017.09.15
**/
object KafkaOffsetManager {

    lazy val log = LogManager.getLogger("KafkaOffsetManager")

    /**
      * 保存每个批次RDD的offset到zk
      * @param zkClient zk客户端
      * @param zkOffsetPath offset存放的zk路径
      * @param rdd 每个批次的RDD
      */
    def setConsumerOffsetsToZK(zkClient: ZkClient, zkOffsetPath : String, rdd : RDD[_]) : Unit = {
        // 将RDD转换为Array[OffsetRange]
        val offsetRanges : Array[OffsetRange] = rdd.asInstanceOf[HasOffsetRanges].offsetRanges
        // 转换每个OffsetRange为存储到zk时的数据格式：分区序号1:偏移量1,分区序号2:偏移量2,分区序号3:偏移量3,……
        val offsetData = offsetRanges.map(offsetRange => s"${offsetRange.partition}:${offsetRange.untilOffset}").mkString(",")
        log.info("保存到zookeeper的offset偏移量 : " + offsetData)
        // 将组装后的offsetData保存到zk上
        ZkUtils.updatePersistentPath(zkClient, zkOffsetPath, offsetData)

    }

    /**
      * 读取zk中保存的offset
      * 若存在新增分区，需要将新增分区也加入进来，并将offset设置为0，保证不会丢失数据
      * @param zkClient zk客户端
      * @param zkOffsetPath offset存放的zk路径
      * @param topic kafka topic名称
      * @return 有则返回分区与对应的偏移量，无则返回none
      */
    def getConsumerOffsetsFromZK(zkClient : ZkClient, zkOffsetPath : String, topic : String) : Option[Map[TopicAndPartition, Long]] = {
        // 从zk上读取偏移量，返回为二元组（偏移量字符串或None，zk元数据）
        val offsetData = ZkUtils.readDataMaybeNull(zkClient, zkOffsetPath)
        // 对返回的二元组第一个元素进行模式匹配
        offsetData._1 match {
            // 若存在值
            case Some(offsetsRanges) =>
                // 获取该topic在zk中的最新分区序号
                val lastest_partitions= ZkUtils.getPartitionsForTopics(zkClient,Seq(topic)).get(topic).get
                log.info("......offset data : " + offsetsRanges + "......")
                // 将之前存于zk的offsetData偏移量字符串做如下几步处理：
                //      1.按逗号切分(split)成数组
                //      2.再将数组的每一个元素按冒号切分(split)得到每个分区(partition)以及对应的偏移量(offset)
                //      3.将每个分区(partition)及分区对应偏移量offset封装于TopicAndParition
                //      4.返回map
                var offsets = offsetsRanges.split(",")
                        .map(s => s.split(":"))
                        .map{
                            case Array(partition, offset) => (TopicAndPartition(topic, partition.toInt) -> offset.toLong)
                        }.toMap
                // 将之前保存的partition分区数量与最新分区数量做对比
                if (offsets.size < lastest_partitions.size) {
                    // 获取之前保存的所有分区的序号
                    val old_partitions = offsets.keys.map(p => p.partition).toArray
                    // 与最新分区做差集，得到本次任务运行新增的kafka分区序号
                    val add_partitions = lastest_partitions.diff(old_partitions)
                    // 若存在新增kafka分区，则需要将新增分区的偏移量设置为0，并加入进来，保证数据不丢失
                    if (add_partitions.size > 0) {
                        log.info("发现kafka新增分区 : " + add_partitions.mkString(","))
                        add_partitions.foreach(partitionId => {
                            offsets += (TopicAndPartition(topic, partitionId) -> 0L)
                            log.info("新增分区id : " + partitionId + "添加成功……")
                        })
                    }
                } else {
                    log.info("没有发现新增的kafka分区: " + lastest_partitions.mkString(","))
                }
                Some(offsets)
            // 若不存在值，则返回None
            case None =>
                None
        }
    }
}
